from face68 import Face68, Response
from body_seg import BodySeg, NeckBodySeg, ClothBodySeg, HairBodySeg, BodySegModel
from utils.ps_canvas import PSCanvasLayer, PSTools
from PP_Matting import PPMatting
import os
import json
import time
from loguru import logger


class Solution:
    def __init__(self, body_seg_model:BodySegModel, face_model:Face68):
        self.need_rm_filenames = []
        self.body_seg_model = body_seg_model
        self.face_model = face_model

    def remove_bg(self):
        pass

    def rm_need_filename(self):
        logger.warning(f"将清除临时文件: {self.need_rm_filenames}")
        for filename in self.need_rm_filenames:
            os.system(f"rm -rf {filename}")


class GenerateTemplateInfo(Solution):
    pixels_to_expand = 8

    def __init__(self, body_seg_model:BodySegModel, face_model:Face68):
        super().__init__(body_seg_model, face_model)
        self.body_model_cloth = ClothBodySeg(model=self.body_seg_model)
        self.body_model_neck = NeckBodySeg(model=self.body_seg_model)

    def _run_(self, image_path, save_dir):
        os.makedirs(save_dir, exist_ok=True)
        # 获取关键点信息
        result = self.face_model.get_keypoints_info(
            image_path=image_path,
            save_dir=save_dir)

        self.need_rm_filenames.append(result.image_path)
        # 记录关键人脸信息
        logger.info(f"message: {result.dict()}")
        with open(save_dir+"/info.json", "w") as f:
            f.write(json.dumps(result.dict()))

        # 提取衣服,并修复孤立点
        mask_filename = self.body_model_cloth.get_mask(
            image_path=image_path, save_dir=save_dir)
        PSTools.mask_remove_isolated_pixels(
            mask_filename, out_path=mask_filename)  # 修复孤立点
        cloth_filename = PSTools.compare_mask_and_origin(
            image_path, mask_filename, save_dir)
        self.need_rm_filenames.append(cloth_filename)
        self.need_rm_filenames.append(mask_filename)
        os.system(f"mv {cloth_filename} {save_dir+'/cloth.png'}")

        # 提取脖子蒙版
        mask_filename = self.body_model_neck.get_mask(
            image_path=image_path,
            save_dir=save_dir)
        neck_filename = PSTools.compare_mask_and_origin(
            image_path, mask_filename, save_dir)  # 生成原始脖子
        os.system(f"mv {neck_filename} {save_dir+'/neck.png'}")
        PSTools.mask_expander(mask_filename, out_path=mask_filename,
                              pixels_to_expand=self.pixels_to_expand)  # 将脖子蒙版扩充,得到inpaint的蒙版
        PSTools.mask_bitwise_not(mask_filename, out_path=mask_filename)
        os.system(f"mv {mask_filename} {save_dir+'/neck-mask.png'}")
        self.need_rm_filenames.append(mask_filename)

    def __call__(self, image_path, save_dir):
        try:
            self._run_(image_path, save_dir)
        except Exception as e:
            logger.exception(str(e))
        finally:
            pass


class GenerateCanvas(Solution):
    canvas_height = 680
    canvas_width = 512
    head_move_pixel = 4
    pixels_to_expand = 10

    def __init__(self, body_seg_model, face_model):
        super().__init__(body_seg_model, face_model)
        self.body_model = BodySeg(model=self.body_seg_model)

    def _get_template_info_(self, path) -> Response:
        with open(path+"/info.json", "r") as f:
            info = json.loads(f.read())

        return Response(**info)

    def _run_(self, image_path, template_dir, save_dir):
        os.makedirs(save_dir, exist_ok=True)
        # 先获取模板信息，得到衣服图层
        template_info: Response = self._get_template_info_(template_dir)
        template_x, template_y = template_info.point_8

        # 第一步先解析人脸信息，将人脸摆正，并根据模板进行缩放人脸，方便直接在图层上直接进行移动
        result = self.face_model.get_keypoints_info(
            image_path=image_path,
            save_dir=save_dir, target_face_width=min(template_info.face_width, 300))

        head_x, head_y = result.point_8
        self.need_rm_filenames.append(result.image_path)

        # 第二步，将摆正后的人脸分割出来,得到人脸图层
        mask_filename = self.body_model.get_mask(
            image_path=result.image_path,
            save_dir=save_dir)
        self.need_rm_filenames.append(mask_filename)
        PSTools.find_max_mask(
            mask_filename, out_path=mask_filename, bitwise_not=True)  # 找最大的脸
        PSTools.mask_remove_isolated_pixels(
            mask_filename, out_path=mask_filename)  # 修复孤立点
        filename = PSTools.compare_mask_and_origin(
            result.image_path, mask_filename, save_dir)  # mask与原像素差值
        self.need_rm_filenames.append(filename)
        # 构建图层
        ps = PSCanvasLayer(canvas_height=self.canvas_height,
                           canvas_width=self.canvas_width)
        ps.add_layer(template_dir+"/neck.png", (0, 0))  # 脖子图层
        ps.add_layer(filename, (template_x-head_x, template_y-head_y))  # 人脸图层
        ps.add_layer(template_dir+"/cloth.png", (0, 0))  # 衣服图层
        filename = ps.merge_layers(save_dir=save_dir)
        return filename

    def __call__(self, image_path, template_dir, save_dir):
        try:
            return self._run_(image_path, template_dir, save_dir)
        except Exception as e:
            logger.exception(str(e))
        finally:
            pass


class PortraitSolution():

    bg_color = "#FFFFFF"

    def __init__(self, body_seg_model:BodySegModel, face_model:Face68, pp_matting_model:PPMatting):
        self.template_client = GenerateTemplateInfo(
            body_seg_model, face_model)
        self.canvas_client = GenerateCanvas(body_seg_model, face_model)
        self.pp_matting_model = pp_matting_model

    def __call__(self, image_path, template_path, save_dir):
        result_type = "Replace BG"
        image_path = self.pp_matting_model(image_path=image_path, save_dir=save_dir,
                                           result_type=result_type, bg_color=self.bg_color, algorithm="ppmatting-512")
        # 模板制作
        time_before = time.time()
        self.template_client(image_path=template_path, save_dir=save_dir)
        time_after = time.time()
        logger.info(f"制作模板总时间: {time_after-time_before}")

        # 开始构建图层
        time_before = time.time()
        filename = self.canvas_client(
            image_path=image_path, template_dir=save_dir, save_dir=save_dir)
        time_after = time.time()
        logger.info(f"图层构建总时间: {time_after-time_before}")
        return filename
